#!/usr/bin/env python
# -*- coding:utf-8 -*-
import pandas as pd
import requests
import time
import json
import pymysql
from datetime import datetime, timedelta, date

def queryGS(difyheaders,dataset_id,document_id,this_date,period):

    # 连接到MySQL数据库 - 使用pymysql方式
    cnx = pymysql.connect(
        user='dolphinadmin',
        password='XjYkRXwjmRRmMmXD',
        host='localhost',
        database='dolphinscheduler',
        charset='utf8mb4'
    )
    cursor = cnx.cursor()
    # 获取今天的日期
    today = this_date
    # 根据period设置时间范围
    if period == "早上-":
        start_time = f"{today} 00:00:00"
        end_time = f"{today} 09:00:00"
    elif period == "中午-":
        start_time = f"{today} 09:00:00"
        end_time = f"{today} 17:00:00"
    elif period == "晚上-":
        start_time = f"{today} 17:00:00"
        end_time = f"{today} 23:59:59"
    else:
        start_time = f"{today} 09:00:00"
        end_time = f"{today} 17:00:00"


    platform_map = {
        "ganggu": "港股",
        "usstock": "美股",
        "astock": "A股",
        "yaowen": "宏观要闻"
    }
    platform_map_stype = {}
    for type in platform_map.keys():
        query = f"SELECT pull_date,response FROM bond_pull_board_data WHERE param = '{type}' AND pull_date >= '{start_time}' AND pull_date < '{end_time}';"
        # 执行查询
        cursor.execute(query)
        results = cursor.fetchall()
        hot_spot_data = [
            {
                "content": row[1],
                "updateTime": row[0].strftime('%Y-%m-%d %H:%M:%S')
            }
            for row in results
        ]
        platform_map_stype.setdefault(type, hot_spot_data)
    cursor.close()
    cnx.close()

    for type in platform_map.keys():
        platform_name = platform_map.get(type, f'未知平台({type})')
        segmentcontent = {
            "keyword": str(this_date)+str(period)+str(platform_name),
            "content": platform_map_stype.get(type, [])
        }
        # 将内容转换为字符串格式
        content_str = json.dumps(segmentcontent, ensure_ascii=False) if segmentcontent else ""
        segment = {
            "content": content_str,
            "keywords": [this_date+str(platform_name)]
        }
        segments_data = {
            "segments": [segment]
        }

        url = "http://localhost:8088/v1/datasets/"+dataset_id+"/documents/"+document_id+"/segments"
        post_with_detailed_retry(url, difyheaders, segments_data)
    return "ok"

def update_document_segments(pull_date,dataset_id, api_key):


    if pull_date=="now":
        now = datetime.now()
        this_date = str(now.strftime('%Y-%m-%d'))
        hour = now.hour
    else:
        this_date = str(pull_date)[:10]
        hour = datetime.strptime(pull_date, '%Y-%m-%d %H:%M:%S').hour
    if hour >17 :
        period = "晚上-"
    elif hour >9 :
        period = "中午-"
    else:
        period = "早上-"

    this_date_srt = str(this_date)+str(period)
    print(f"当前日期: {this_date_srt}")

    url = "http://localhost:8088/v1/datasets/"+dataset_id+"/document/create-by-text"
    headers = {
        'Authorization': f'Bearer {api_key}',
        'Content-Type': 'application/json'
    }
    payload = {
        "name": this_date_srt+"财经数据",
        "text": "",
        "indexing_technique": "high_quality",
        "doc_form": "text_model",
        "process_rule": {"mode": "automatic"}
    }
    response = requests.post(url, headers=headers, data=json.dumps(payload))
    json_data = response.json()
    print(f"[{datetime.now()}] 创建文档返回 数据: {json_data}")

    document_id =json_data['document']['id']
    print(f"document_id: {document_id}")
    time.sleep(5)
    newrankXhs_data = queryGS(headers,dataset_id,document_id,this_date,period)
    return "全部完成"


def post_with_detailed_retry(url, headers, text_data):
    print(f"[{datetime.now()}] 开始发送POST请求到: {url}")
    try:
        responsesegments = requests.post(
            url,
            headers=headers,
            data=json.dumps(text_data)
        )
        # 如果不是404，成功返回
        if responsesegments.status_code != 404:
            print(f"[{datetime.now()}] 请求成功!")
    except Exception as e:
        print(f"[{datetime.now()}] 请求segments出现异常: {e}")
    return "ok"


if __name__ == "__main__":
    dataset_id = '${dataset_id}'
    api_key = '${api_key}'
    pull_date = '${pull_date}'
    print("开始 getGS")
    response_data = update_document_segments(pull_date,dataset_id, api_key)
    print("结束 getGS")